;;; guile-openai --- An OpenAI API client for Guile
;;; Copyright © 2023 Andrew Whatson <whatson@tailcall.au>
;;;
;;; This file is part of guile-openai.
;;;
;;; guile-openai is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU Affero General Public License as
;;; published by the Free Software Foundation, either version 3 of the
;;; License, or (at your option) any later version.
;;;
;;; guile-openai is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; Affero General Public License for more details.
;;;
;;; You should have received a copy of the GNU Affero General Public
;;; License along with guile-openai.  If not, see
;;; <https://www.gnu.org/licenses/>.

(define-module (openai api completion)
  #:use-module (openai client)
  #:use-module (json record)
  #:use-module (srfi srfi-41)
  #:export (make-completion-request
            json->completion-request
            completion-request->json
            completion-request?
            completion-request-model
            completion-request-prompt
            completion-request-suffix
            completion-request-max-tokens
            completion-request-temperature
            completion-request-top-p
            completion-request-n
            completion-request-stream?
            completion-request-logprobs
            completion-request-echo
            completion-request-stop
            completion-request-presence-penalty
            completion-request-frequency-penalty
            completion-request-best-of
            completion-request-logit-bias
            completion-request-user

            make-completion-response
            json->completion-response
            completion-response->json
            completion-response?
            completion-response-id
            completion-response-object
            completion-response-created
            completion-response-model
            completion-response-choices
            completion-response-usage

            make-completion-choice
            json->completion-choice
            completion-choice->json
            completion-choice?
            completion-choice-text
            completion-choice-index
            completion-choice-logprobs
            completion-choice-finish-reason

            make-completion-logprobs
            json->completion-logprobs
            completion-logprobs->json
            completion-logprobs?
            completion-logprobs-tokens
            completion-logprobs-token-logprobs
            completion-logprobs-top-logprobs
            completion-logprobs-completion-offset

            make-completion-usage
            json->completion-usage
            completion-usage->json
            completion-usage?
            completion-usage-prompt-tokens
            completion-usage-completion-tokens
            completion-usage-total-tokens

            send-completion-request))

;; See https://platform.openai.com/docs/api-reference/completions

(define-json-type <completion-request>
  (model)
  (prompt)
  ;; XXX: suffix causes an API error
  (suffix)
  (max-tokens        "max_tokens")
  (temperature)
  (top-p             "top_p")
  (n)
  (stream?           "stream")
  (logprobs)
  (echo)
  (stop)
  (presence-penalty  "presence_penalty")
  (frequency-penalty "frequency_penalty")
  (best-of           "best_of")
  (logit-bias        "logit_bias")
  (user))

(define-json-type <completion-response>
  (id)
  (object)
  (created)
  (model)
  (choices "choices" #(<completion-choice>))
  (usage   "usage"   <completion-usage>))

(define-json-type <completion-choice>
  (text)
  (index)
  (logprobs)
  ;; XXX: logprobs can be null or array
  ;; (logprobs "logprobs" #(<completion-logprobs>))
  (finish-reason "finish_reason"))

(define-json-type <completion-logprobs>
  (tokens)
  (token-logprobs "token_logprobs")
  (top-logprobs   "top_logprobs")
  (text-offset    "text_offset"))

(define-json-type <completion-usage>
  (prompt-tokens     "prompt_tokens")
  (completion-tokens "completion_tokens")
  (total-tokens      "total_tokens"))

(define (send-completion-request request)
  (let ((result (openai-post-json "/v1/completions"
                                  (completion-request->json request))))
    (if (stream? result)
        (stream-map json->completion-response result)
        (json->completion-response result))))
